import pandas as pd
import logging
from typing import List, Dict, Optional
from datetime import datetime, timedelta
from src.database.data_manager import DataManager
from decimal import Decimal

logger = logging.getLogger(__name__)

class DragonRepairStrategy:
    """
    龙修复策略
    股票最后三个交易日A1,A2,A3的数据，
    满足以下两种情况的其中一种则可以选中股票：
    一、A2开盘涨停,但是收盘不是涨停,而且A3低开高走,最终收盘价低于A2收盘价格。记录为情况1;
    二、A3开盘涨停,但是收盘不是涨停。记录为情况2
    """
    
    def __init__(self, strategy_manager=None):
        self.strategy_manager = strategy_manager
        self.data_manager = DataManager.get_instance()
        # 普通股票池
        self.stock_pool = []
        # 初始化股票池
        self._init_stock_pool()
        
    def _init_stock_pool(self):
        try:
            # 直接从DataManager获取普通股票列表
            self.stock_pool = self.strategy_manager.get_normal_stocks()
            if len(self.stock_pool) <= 0:
                logger.warning("没有普通股票数据")
                return
            logger.info(f"初始化股票池完成，共{len(self.stock_pool)}只股票")
            
        except Exception as e:
            logger.error(f"初始化股票池失败: {str(e)}")
            self.stock_pool = []

    def run(self) -> List[Dict]:
        """运行策略"""
        try:
            if len(self.stock_pool) <= 0:
                self._init_stock_pool()
            
            selected_stocks = []
            
            # 遍历DataFrame的每一行
            for _, stock in self.stock_pool.iterrows():
                # 从Series中获取股票代码和名称
                code = stock['code']
                name = stock['name']
                
                if pd.isna(code) or pd.isna(name):
                    continue
                
                # 直接从PostgreSQL获取最近3个交易日的数据
                df = self.data_manager.postgres.query_stock_prices(
                    symbol=code,
                    limit=3,
                    order_by='date DESC'
                )
                
                if len(df) < 3:
                    continue
                    
                # 按日期正序排列，使A1、A2、A3分别为最近三个交易日
                df = df.sort_values('date')
                A1 = df.iloc[0]
                A2 = df.iloc[1]
                A3 = df.iloc[2]
                
                # 计算涨停价（转换Decimal类型为float进行计算）
                limit_up_A2 = round(float(A1['close']) * 1.1, 2)
                limit_up_A3 = round(float(A2['close']) * 1.1, 2)
                A2_open = float(A2['open'])
                A2_close = float(A2['close'])
                A3_open = float(A3['open'])
                A3_close = float(A3['close'])
                
                # 情况1：A2开盘涨停，但收盘不涨停，且A3低开高走，收盘低于A2收盘
                condition1 = (
                    (abs(A2_open - limit_up_A2) < 0.02 or A2_open - limit_up_A2 > 0)   # A2开盘涨停
                    and A2_close < limit_up_A2 - 0.03 # A2收盘未涨停
                    and  A3_open < A3_close  # A3低开高走
                    and  A3_close < A2_close  # A3收盘低于A2收盘
                )
                
                # 情况2：A3开盘涨停但收盘不涨停
                condition2 = (
                    (abs(A3_open - limit_up_A3) < 0.02 or A3_open - limit_up_A3 >= 0)  # A3开盘涨停
                     and A3_close < limit_up_A3 - 0.02  # A3收盘未涨停
                )
                
                if condition1 or condition2:
                    stock_info = {
                        'code': code,
                        'name': name,
                        'date': A3['date'].strftime('%Y%m%d'),
                        'reason': '情况1' if condition1 else '情况2',
                        'price_data': {
                            'A1': {
                                'date': A1['date'].strftime('%Y%m%d'),
                                'open': float(A1['open']),
                                'close': float(A1['close'])
                            },
                            'A2': {
                                'date': A2['date'].strftime('%Y%m%d'),
                                'open': A2_open,
                                'close': A2_close
                            },
                            'A3': {
                                'date': A3['date'].strftime('%Y%m%d'),
                                'open': A3_open,
                                'close': A3_close
                            }
                        }
                    }
                    selected_stocks.append(stock_info)
            
            logger.info(f"策略选出{len(selected_stocks)}只股票")
            return selected_stocks
            
        except Exception as e:
            logger.error(f"策略运行失败: {str(e)}")
            return []
